from exper.Experiment import Experiment
import torch
import os
import math
import numpy as np
import matplotlib.pyplot as plt
import math
import seaborn as sns

class ExpTransfer(Experiment):

    def __init__(self, config, dataer, neter):

        super(ExpTransfer, self).__init__('Transfer', config)
        self.dataer = dataer
        self.neter = neter
        self.neter.Reload_network(model_arch=self.config['transfer_model_arch'])
  
    def Case_transfer_MPack(self, seed, repeat_times=3, default_num=500, way='MPack', atk=None):

        ASR_list = []
        avg_Query_list = []
        median_Query_list = []

        for step in range(repeat_times):

            loader = self.dataer.get_shuffle_loader(seed=seed+step, isTrain=False, batch_size=1)
            succ_number = 0
            query_buffer = 0
            single_sample_query_list = []

            count = 0
            total = 0

            for (image, label) in loader:
                
                if count == default_num:
                    break
                total += 1
                test_image = image.to(self.neter.device)
                test_label = label.to(self.neter.device)
                output = self.neter.net(test_image)
                _, pre = torch.max(output.data, 1)
                if pre != test_label:
                    print(pre)
                    print(test_label)
                    print('Pass the error sample')
                    continue

                if way == 'SMPack':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.untarget_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                        batch=self.config['base_batch'],
                                        size=self.config['base_size'],
                                        mode=self.config['base_mode'],
                                    )
                    if flag:
                        succ_number += 1         
                
                elif way == 'GA':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.untarget_GA_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                    )
                    if flag:
                        succ_number += 1     
                   
                elif way == 'MPack':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.untarget_MPack_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                        batch=self.config['base_batch'],
                                        size=self.config['base_size'],
                                        mode=self.config['base_mode'],
                                    )
                    if flag:
                        succ_number += 1   

                elif way == 'FNS':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.FNS_attack(
                                        images=image,
                                        labels=label,
                                    )
                    if flag:
                        succ_number += 1 

                elif way == 'OnePixel':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1   

                elif way == 'SimBA':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.SimBA_attack(image, label[0], max_queries=10000, epsilon=self.config['epsilon'], targeted=False)

                    if flag:
                        succ_number += 1
                        
                elif way == 'FGSM':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1  
                
                elif way == 'PGD':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1  

                elif way == 'ZOO':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image, _ = atk.attack(image, label)
                    adv_image = torch.from_numpy(adv_image).unsqueeze(0)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1  

                elif way == 'Test':
                    pass
                else:
                    raise Exception('No such compared method.')
                print('Total : {}'.format(total))
                ## query counts
                print('='*50)
                print('Qurey for {}-th samples : {}'.format(count, self.neter.query - query_buffer))
                count += 1
                print('Total Query: {}'.format(self.neter.query))
                print('Last Total Query: {}'.format(query_buffer))
                print('='*50)
                single_sample_query_list.append(self.neter.query - query_buffer)
                query_buffer = self.neter.query

            print('Way {} : The Epoch {}-th Untarget Attack ASR (Randomly sample {} samples) : {:.2f}%'.format(way, step, default_num, succ_number / default_num * 100))
            
            ASR_list.append(succ_number / default_num) 
            avg_Query_list.append(self.neter.query / default_num)
            
            # median
            single_sample_query_list.sort()
            list_len = len(single_sample_query_list)
            if list_len % 2 == 0:
                median_Query_list.append((single_sample_query_list[int(list_len / 2) - 1] + single_sample_query_list[int(list_len / 2)]) / 2)
            else:
                median_Query_list.append(single_sample_query_list[int(list_len / 2)])

            self.neter.clear_query()

        case_1_path = self.generate_name(case='Untarget-{}'.format(self.config['model_arch']))

        # for case 1:
        if os.path.exists(case_1_path) == False:
            os.makedirs(case_1_path)
        if os.path.exists(os.path.join(case_1_path, 'result.pt')) == False:
            case1_dicter = {}
            case1_dicter['{}'.format(self.config['transfer_model_arch'])] = [np.mean(ASR_list), np.var(ASR_list)]
        else:
            case1_dicter = torch.load(os.path.join(case_1_path, 'result.pt'))
            case1_dicter['{}'.format(self.config['transfer_model_arch'])] = [np.mean(ASR_list), np.var(ASR_list)]
        torch.save(case1_dicter, os.path.join(case_1_path, 'result.pt'))

    def Case_transfer_target_MPack(self, seed, repeat_times=3, default_num=500, way='MPack', atk=None):

        ASR_list = []
        avg_Query_list = []
        median_Query_list = []

        for step in range(repeat_times):

            loader = self.dataer.get_shuffle_loader(seed=seed+step, isTrain=False, batch_size=1)
            succ_number = 0
            query_buffer = 0
            single_sample_query_list = []

            count = 0
            total = 0
            
            ##target attack mode: attacking the target class: [(label + 1) mod total_class_num] 
            

            for (image, label) in loader:
                
                if count == default_num:
                    break

                total += 1

                ## 确保该原始样本被正确分类
                test_image = image.to(self.neter.device)
                test_label = label.to(self.neter.device)
                output = self.neter.net(test_image)
                _, pre = torch.max(output.data, 1)
                if pre != test_label:
                    print(pre)
                    print(test_label)
                    print('Pass the error sample')
                    continue    
                
                ## implement of attacking the (label + 1) % total_class_num
                target_class = (label[0] + 1) % self.dataer.class_num

                if way == 'SMPack':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    flag, adv_image = atk.target_attck(
                                        origin_sample=image,
                                        source_class=label[0],
                                        target_class=target_class,
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                        batch=self.config['base_batch'],
                                        size=self.config['base_size'],
                                        mode=self.config['base_mode'],
                                    )
                    if flag:
                        succ_number += 1         
                
                elif way == 'GA':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    flag, adv_image = atk.target_GA_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        target_class=target_class,
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                    )
                    if flag:
                        succ_number += 1     
                   
                elif way == 'MPack':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    flag, adv_image = atk.target_MPack_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        target_class=target_class,
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                        batch=self.config['base_batch'],
                                        size=self.config['base_size'],
                                        mode=self.config['base_mode'],
                                    )
                    if flag:
                        succ_number += 1   

                elif way == 'FNS':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    atk.set_attack_mode(isTarget=True, class_index=target_class)
                    flag, adv_image = atk.FNS_attack(
                                        images=image,
                                        labels=label,
                                    )
                    if flag:
                        succ_number += 1 

                elif way == 'OnePixel':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                                       
                    # change into target attack mode
                    atk.set_mode_targeted_by_function(target_map_function=lambda image, label: (label + 1) % self.dataer.class_num)
                    
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre[0] == target_class:                   
                        succ_number += 1   

                elif way == 'SimBA':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    flag, adv_image = atk.SimBA_attack(image, target_class, max_queries=10000, epsilon=self.config['epsilon'], targeted=True)

                    if flag:
                        succ_number += 1
                        
                elif way == 'FGSM':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
            
                    # change into target attack mode
                    atk.set_mode_targeted_by_function(target_map_function=lambda image, label: (label + 1) % self.dataer.class_num)
    
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre[0] == target_class:                   
                        succ_number += 1  

                elif way == 'PGD':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
            
                    # change into target attack mode
                    atk.set_mode_targeted_by_function(target_map_function=lambda image, label: (label + 1) % self.dataer.class_num)
    
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre[0] == target_class:                   
                        succ_number += 1  

                elif way == 'ZOO':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    adv_image, _ = atk.attack(image, label)
                    adv_image = torch.from_numpy(adv_image).unsqueeze(0)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1  

                elif way == 'Test':
                    pass

                else:
                    raise Exception('No such compared method.')

                print('Total : {}'.format(total))
                ## query counts
                print('='*50)
                print('Qurey for {}-th samples : {}'.format(count, self.neter.query - query_buffer))
                count += 1
                print('Total Query: {}'.format(self.neter.query))
                print('Last Total Query: {}'.format(query_buffer))
                print('='*50)
                single_sample_query_list.append(self.neter.query - query_buffer)
                query_buffer = self.neter.query

            print('Way {} : The Epoch {}-th Target Attack ASR (Randomly sample {} samples) : {:.2f}%'.format(way, step, default_num, succ_number / default_num * 100))
            
            ASR_list.append(succ_number / default_num) 
            avg_Query_list.append(self.neter.query / default_num)
            
            # median
            single_sample_query_list.sort()
            list_len = len(single_sample_query_list)
            if list_len % 2 == 0:
                median_Query_list.append((single_sample_query_list[int(list_len / 2) - 1] + single_sample_query_list[int(list_len / 2)]) / 2)
            else:
                median_Query_list.append(single_sample_query_list[int(list_len / 2)])

            self.neter.clear_query()

        case_1_path = self.generate_name(case='Target-{}'.format(self.config['model_arch']))

        # for case 1:
        if os.path.exists(case_1_path) == False:
            os.makedirs(case_1_path)
        if os.path.exists(os.path.join(case_1_path, 'result.pt')) == False:
            case1_dicter = {}
            case1_dicter['{}'.format(self.config['transfer_model_arch'])] = [np.mean(ASR_list), np.var(ASR_list)]
        else:
            case1_dicter = torch.load(os.path.join(case_1_path, 'result.pt'))
            case1_dicter['{}'.format(self.config['transfer_model_arch'])] = [np.mean(ASR_list), np.var(ASR_list)]
        torch.save(case1_dicter, os.path.join(case_1_path, 'result.pt'))

    def Case_speed(self, seed, repeat_times=3, default_num=100, way='MPack', atk=None, total_query=200000):

        ## s设定总的查询次数，以及单张样本的查询上限(注意，如果算法在总的次数内完成了目标，要妥善处理好后续，code层面)。
        ## 单张图片的的最大访问次数是10000次，不成功就下一张
        ## 比较方法是GA SMPack-VGG16 SMPack-ResNet
        ## 给出存储方式
        ## 以10次查询为单位进行打点

        ASR_list = []
        Query_list = []

        for step in range(repeat_times):

            sub_ASR_list = []
            sub_Query_list = []

            loader = self.dataer.get_shuffle_loader(seed=seed+step, isTrain=False, batch_size=1)
            succ_number = 0
            query_buffer = 0
            single_sample_query_list = []

            count = 0
            total = 0

            for (image, label) in loader:
                
                if count == default_num:
                    break
                total += 1
                test_image = image.to(self.neter.device)
                test_label = label.to(self.neter.device)
                output = self.neter.net(test_image)
                _, pre = torch.max(output.data, 1)
                if pre != test_label:
                    print(pre)
                    print(test_label)
                    print('Pass the error sample')
                    continue

                if way == 'SMPack':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.untarget_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                        batch=self.config['base_batch'],
                                        size=self.config['base_size'],
                                        mode=self.config['base_mode'],
                                    )
                    if flag:
                        succ_number += 1         
                
                elif way == 'GA':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.untarget_GA_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                    )
                    if flag:
                        succ_number += 1     
                   
                elif way == 'MPack':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.untarget_MPack_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                        batch=self.config['base_batch'],
                                        size=self.config['base_size'],
                                        mode=self.config['base_mode'],
                                    )
                    if flag:
                        succ_number += 1   

                elif way == 'FNS':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.FNS_attack(
                                        images=image,
                                        labels=label,
                                    )
                    if flag:
                        succ_number += 1 

                elif way == 'OnePixel':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1   

                elif way == 'SimBA':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    flag, adv_image = atk.SimBA_attack(image, label[0], max_queries=10000, epsilon=self.config['epsilon'], targeted=False)

                    if flag:
                        succ_number += 1
                        
                elif way == 'FGSM':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1  
                
                elif way == 'PGD':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1  

                elif way == 'ZOO':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample'.format(step, count))
                    adv_image, _ = atk.attack(image, label)
                    adv_image = torch.from_numpy(adv_image).unsqueeze(0)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1  

                elif way == 'Test':
                    pass

                else:
                    raise Exception('No such compared method.')

                print('Total : {}'.format(total))
                ## query counts
                print('='*50)
                print('Qurey for {}-th samples : {}'.format(count, self.neter.query - query_buffer))
                count += 1
                print('Total Query: {}'.format(self.neter.query))
                print('Last Total Query: {}'.format(query_buffer))
                print('='*50)

                for query_count in range(query_buffer, self.neter.query, 10):

                    if succ_number == 0:
                        sub_ASR_list.append(0.0)
                        sub_Query_list.append(query_count)
                    else:
                        sub_ASR_list.append((succ_number - 1) / default_num)
                        sub_Query_list.append(query_count)
                
                query_buffer = self.neter.query

            print('Way {} : The Epoch {}-th Untarget Attack ASR (Randomly sample {} samples) : {:.2f}%'.format(way, step, default_num, succ_number / default_num * 100))
            
            ## extension lenth of the sub_list
            ## TODO::
            real_lenth = len(sub_ASR_list)
            for query_count in range(real_lenth * 10, total_query, 10):
                sub_ASR_list.append(succ_number / default_num)
                sub_Query_list.append(query_count)
            
            if len(sub_ASR_list) != int(total_query / 10) or len(sub_Query_list) != int(total_query / 10):
                raise ValueError('The array list lenth is {} and {}, not {} !'.format(len(sub_ASR_list), len(sub_Query_list), int(total_query / 10)))

            ASR_list.append(sub_ASR_list)
            Query_list.append(sub_Query_list)

            self.neter.clear_query()

        ASR_array = np.array(ASR_list)
        Query_array = np.array(Query_list)

        asr = np.mean(ASR_array, axis=0)
        query = np.mean(Query_array, axis=0)

        ## TODO:: 本地重写generate_name方法
        case_2_path = self.generate_name(case='Speed-{}->{}'.format(self.config['model_arch'], self.config['transfer_model_arch']))

        # for case 1:
        if os.path.exists(case_2_path) == False:
            os.makedirs(case_2_path)
        if os.path.exists(os.path.join(case_2_path, 'result.pt')) == False:
            case1_dicter = {}
            case1_dicter[way] = {'ASR': asr, 'Query': query}
        else:
            case1_dicter = torch.load(os.path.join(case_2_path, 'result.pt'))
            case1_dicter[way] = {'ASR': asr, 'Query': query}
        torch.save(case1_dicter, os.path.join(case_2_path, 'result.pt'))

    def Case_target_speed(self, seed, repeat_times=3, default_num=100, way='MPack', atk=None, total_query=200000):

        ## s设定总的查询次数，以及单张样本的查询上限(注意，如果算法在总的次数内完成了目标，要妥善处理好后续，code层面)。
        ## 单张图片的的最大访问次数是10000次，不成功就下一张
        ## 比较方法是GA SMPack-VGG16 SMPack-ResNet
        ## 给出存储方式
        ## 以10次查询为单位进行打点

        ASR_list = []
        Query_list = []

        for step in range(repeat_times):

            sub_ASR_list = []
            sub_Query_list = []

            loader = self.dataer.get_shuffle_loader(seed=seed+step, isTrain=False, batch_size=1)
            succ_number = 0
            query_buffer = 0

            count = 0
            total = 0

            for (image, label) in loader:
                
                if count == default_num:
                    break

                total += 1

                ## 确保该原始样本被正确分类
                test_image = image.to(self.neter.device)
                test_label = label.to(self.neter.device)
                output = self.neter.net(test_image)
                _, pre = torch.max(output.data, 1)
                if pre != test_label:
                    print(pre)
                    print(test_label)
                    print('Pass the error sample')
                    continue    
                
                ## implement of attacking the (label + 1) % total_class_num
                target_class = (label[0] + 1) % self.dataer.class_num

                if way == 'SMPack':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    flag, adv_image = atk.target_attck(
                                        origin_sample=image,
                                        source_class=label[0],
                                        target_class=target_class,
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                        batch=self.config['base_batch'],
                                        size=self.config['base_size'],
                                        mode=self.config['base_mode'],
                                    )
                    if flag:
                        succ_number += 1         
                
                elif way == 'GA':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    flag, adv_image = atk.target_GA_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        target_class=target_class,
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                    )
                    if flag:
                        succ_number += 1     
                   
                elif way == 'MPack':
                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    flag, adv_image = atk.target_MPack_attack(
                                        origin_sample=image,
                                        source_class=label[0],
                                        target_class=target_class,
                                        epsilon=self.config['epsilon'],
                                        norm_type=self.config['attck_norm'],
                                        method=self.config['attck_method'],
                                        batch=self.config['base_batch'],
                                        size=self.config['base_size'],
                                        mode=self.config['base_mode'],
                                    )
                    if flag:
                        succ_number += 1   

                elif way == 'FNS':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    atk.set_attack_mode(isTarget=True, class_index=target_class)
                    flag, adv_image = atk.FNS_attack(
                                        images=image,
                                        labels=label,
                                    )
                    if flag:
                        succ_number += 1 

                elif way == 'OnePixel':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                                       
                    # change into target attack mode
                    atk.set_mode_targeted_by_function(target_map_function=lambda image, label: (label + 1) % self.dataer.class_num)
                    
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre[0] == target_class:                   
                        succ_number += 1   

                elif way == 'SimBA':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    flag, adv_image = atk.SimBA_attack(image, target_class, max_queries=10000, epsilon=self.config['epsilon'], targeted=True)

                    if flag:
                        succ_number += 1
                        
                elif way == 'FGSM':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
            
                    # change into target attack mode
                    atk.set_mode_targeted_by_function(target_map_function=lambda image, label: (label + 1) % self.dataer.class_num)
    
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre[0] == target_class:                   
                        succ_number += 1  

                elif way == 'PGD':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
            
                    # change into target attack mode
                    atk.set_mode_targeted_by_function(target_map_function=lambda image, label: (label + 1) % self.dataer.class_num)
    
                    adv_image = atk(image, label)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre[0] == target_class:                   
                        succ_number += 1  

                elif way == 'ZOO':

                    if atk == None:
                        raise ValueError
                    print('Epoch {}, Attacking {}-th sample, source class: {}, target class: {}'.format(step, count, label[0], target_class))
                    adv_image, _ = atk.attack(image, label)
                    adv_image = torch.from_numpy(adv_image).unsqueeze(0)

                    adv_image = adv_image.to(self.neter.device)
                    output = self.neter.net(adv_image)
                    _, pre = torch.max(output.data, 1)

                    if pre != test_label:                   
                        succ_number += 1  

                elif way == 'Test':
                    pass

                else:
                    raise Exception('No such compared method.')

                print('Total : {}'.format(total))
                ## query counts
                print('='*50)
                print('Qurey for {}-th samples : {}'.format(count, self.neter.query - query_buffer))
                count += 1
                print('Total Query: {}'.format(self.neter.query))
                print('Last Total Query: {}'.format(query_buffer))
                print('='*50)

                for query_count in range(query_buffer, self.neter.query, 10):

                    if succ_number == 0:
                        sub_ASR_list.append(0.0)
                        sub_Query_list.append(query_count)
                    else:
                        sub_ASR_list.append((succ_number - 1) / default_num)
                        sub_Query_list.append(query_count)
                
                query_buffer = self.neter.query

            print('Way {} : The Epoch {}-th Untarget Attack ASR (Randomly sample {} samples) : {:.2f}%'.format(way, step, default_num, succ_number / default_num * 100))
            
            ## extension lenth of the sub_list
            ## TODO::
            real_lenth = len(sub_ASR_list)
            if real_lenth > int(total_query / 10):
                 sub_ASR_list = sub_ASR_list[:int(total_query / 10)]
                 sub_Query_list = sub_Query_list[:int(total_query / 10)]
            else:
                for query_count in range(real_lenth * 10, total_query, 10):
                    sub_ASR_list.append(succ_number / default_num)
                    sub_Query_list.append(query_count)
            
            if len(sub_ASR_list) != int(total_query / 10) or len(sub_Query_list) != int(total_query / 10):
                raise ValueError('The array list lenth is {} and {}, not {} !'.format(len(sub_ASR_list), len(sub_Query_list), int(total_query / 10)))

            ASR_list.append(sub_ASR_list)
            Query_list.append(sub_Query_list)

            self.neter.clear_query()

        ASR_array = np.array(ASR_list)
        Query_array = np.array(Query_list)

        asr = np.mean(ASR_array, axis=0)
        query = np.mean(Query_array, axis=0)

        ## TODO:: 本地重写generate_name方法
        case_2_path = self.generate_name(case='Speed-target-{}->{}'.format(self.config['model_arch'], self.config['transfer_model_arch']))

        # for case 1:
        if os.path.exists(case_2_path) == False:
            os.makedirs(case_2_path)
        if os.path.exists(os.path.join(case_2_path, 'result.pt')) == False:
            case1_dicter = {}
            case1_dicter[way] = {'ASR': asr, 'Query': query}
        else:
            case1_dicter = torch.load(os.path.join(case_2_path, 'result.pt'))
            case1_dicter[way] = {'ASR': asr, 'Query': query}
        torch.save(case1_dicter, os.path.join(case_2_path, 'result.pt'))

    def show_result(self, case='Untarget'):

        if case == 'Untarget' or case == 'Target':
            
            path = os.path.join(self.generate_name(case='{}-{}'.format(case, self.config['model_arch'])), 'result.pt')
            if os.path.exists(path) == False:
                raise Exception('Do not find result in path {}'.format(os.path.join(self.generate_name(case='{}-{}'.format(case, self.config['model_arch'])), 'result.pt')))
            dicter = torch.load(os.path.join(self.generate_name(case='{}-{}'.format(case, self.config['model_arch'])), 'result.pt'))

            print(dicter)

            str = ''
            for key, value in dicter.items():
                str += '&${:.1f}\\pm{:.1f}$ '.format(dicter[key][0] * 100, math.sqrt(dicter[key][1]) * 100)
            str += '\\\\'
            
            print(str)

        elif case == 'Speed':

            path = os.path.join(self.generate_name(case='Speed-{}->{}'.format(self.config['model_arch'], self.config['transfer_model_arch'])), 'result.pt')
            if os.path.exists(path) == False:
                raise Exception('Do not find result in path {}'.format(path))
            dicter = torch.load(path)

            print(dicter.keys())
            for key, value in dicter.items():
                print(key, end=' : ')
                print(len(value))

        elif case == 'Speed-target':

            path = os.path.join(self.generate_name(case='Speed-target-{}->{}'.format(self.config['model_arch'], self.config['transfer_model_arch'])), 'result.pt')
            if os.path.exists(path) == False:
                raise Exception('Do not find result in path {}'.format(path))
            dicter = torch.load(path)

            print(dicter.keys())
            for key, value in dicter.items():
                print(key, end=' : ')
                print(len(value))

        else:

            raise ValueError
        
    def Speed_plot(
            self, 
            total_query=200000,
        ):
            ## 构建源数据集到目标数据集的查询次数到翻转率的对比，
            ## transfer_dataset: 被迁移到的数据集
            ## 绘制三条曲边，分别是SMPack(VGG-16)-->Transfer_dataset, SMPack(ResNet-18)-->Transfer_dataset, GA-->Transfer_dataset
            source_dataset_list = ['VGG16', 'ResNet18']

            transfer_datasets = [
                'VGG16', 'ResNet18',
                'VGG19', 'ResNet34',
                'GoogLeNet', 'NiN',
                'LeNet', 'AlexNet',
            ]

            count = 1

            sns.set_style('whitegrid')
            fig = plt.figure(figsize=(16, 22))

            for transfer_dataset in transfer_datasets:

                x = None
                data_list = []

                ax = plt.subplot(4, 2, count)

                for source_dataset in source_dataset_list:
                    self.config['model_arch'] = source_dataset
                    path = os.path.join(self.generate_name(case='Speed-{}->{}'.format(source_dataset, transfer_dataset)), 'result.pt')
                    if os.path.exists(path) == False:
                        raise Exception('Do not find result in path {}'.format(path))
                    dicter = torch.load(path)
                
                    ## checking
                    for method in ['SMPack', 'GA']:
                        sub_dicter = dicter[method]
                        data = sub_dicter['ASR']
                        query = sub_dicter['Query']
                        if len(data) != int(total_query / 10) or len(query) != int(total_query / 10):
                            raise ValueError('The len of data is {} and  query list is {}, but not {}'.format(len(data), len(query), int(total_query / 10)))
                        
                        x = query
                        data_list.append(data)
                
                plt.title('Transferability-->{}'.format(transfer_dataset), fontsize=22)

                ax.plot(x, data_list[0], label='SMPack(VGG-16)->{}'.format(transfer_dataset), linewidth=3.0)
                ax.plot(x, data_list[2], label='SMPack(ResNet-18)->{}'.format(transfer_dataset), linewidth=3.0)
                ax.plot(x, data_list[1], label='GA->{}'.format(transfer_dataset), linewidth=3.0)
                ## plt.plot(x, data_list[3], label='GA(ResNet-18)->{}'.format(transfer_dataset))

                plt.xticks(fontsize=11)
                plt.yticks(fontsize=13)


                ax.legend(fontsize=16)

                count += 1

            ## plt.subplots_adjust(wspace=0.15, hspace=0.25)
            plt.subplots_adjust(wspace=0.1, hspace=0.3)
            plt.savefig('./img/Transferability_Untarget.pdf', dpi=800, bbox_inches='tight', pad_inches=0.2)
            print('Save Done.')
            plt.close()

    def Speed_target_plot(
            self, 
            total_query=200000,
        ):
            ## 构建源数据集到目标数据集的查询次数到翻转率的对比，
            ## transfer_dataset: 被迁移到的数据集
            ## 绘制三条曲边，分别是SMPack(VGG-16)-->Transfer_dataset, SMPack(ResNet-18)-->Transfer_dataset, GA-->Transfer_dataset
            source_dataset_list = ['VGG16', 'ResNet18']

            transfer_datasets = [
                'VGG16', 'ResNet18',
            ]

            count = 1

            sns.set_style('whitegrid')
            fig = plt.figure(figsize=(16, 22))

            for transfer_dataset in transfer_datasets:

                x = None
                data_list = []

                ax = plt.subplot(4, 2, count)

                for source_dataset in source_dataset_list:
                    self.config['model_arch'] = source_dataset
                    path = os.path.join(self.generate_name(case='Speed-target-{}->{}'.format(source_dataset, transfer_dataset)), 'result.pt')
                    if os.path.exists(path) == False:
                        raise Exception('Do not find result in path {}'.format(path))
                    dicter = torch.load(path)
                
                    ## checking
                    for method in ['SMPack', 'GA']:
                        sub_dicter = dicter[method]
                        data = sub_dicter['ASR']
                        query = sub_dicter['Query']
                        if len(data) != int(total_query / 10) or len(query) != int(total_query / 10):
                            raise ValueError('The len of data is {} and  query list is {}, but not {}'.format(len(data), len(query), int(total_query / 10)))
                        
                        x = query
                        data_list.append(data)
                
                plt.title('Transferability-->{}'.format(transfer_dataset), fontsize=22)

                ax.plot(x, data_list[0], label='SMPack(VGG-16)->{}'.format(transfer_dataset), linewidth=3.0)
                ax.plot(x, data_list[2], label='SMPack(ResNet-18)->{}'.format(transfer_dataset), linewidth=3.0)
                ax.plot(x, data_list[1], label='GA->{}'.format(transfer_dataset), linewidth=3.0)
                # ax.plot(x, data_list[3], label='GA(ResNet-18)->{}'.format(transfer_dataset))

                plt.xticks(fontsize=11)
                plt.yticks(fontsize=13)


                ax.legend(fontsize=16)

                count += 1

            plt.subplots_adjust(wspace=0.1, hspace=0.3)
            plt.savefig('./img/Transferability_Target.pdf', dpi=800, bbox_inches='tight', pad_inches=0.2)
            print('Save Done.')
            plt.close()
